import torch


class VariableMemory(torch.nn.Module):
    def __init__(self, value: torch.Tensor | None = None):
        super().__init__()
        if value is not None:
            self._value = torch.nn.Buffer(value)
        else:
            self._value = torch.nn.UninitializedBuffer()

    def get(self) -> torch.Tensor | None:
        if isinstance(self._value, torch.nn.UninitializedBuffer):
            return None
        return self._value

    def set(self, value: torch.Tensor, *, mask: torch.Tensor | None = None):
        if isinstance(self._value, torch.nn.UninitializedBuffer):
            self._value = value
        if mask is None:
            self._value.copy_(value)
        else:
            self._value.masked_scatter_(mask, value)
